// CSc 522, hw1
// MPI version of red/black grid computation
//Jinyan Huan (jguan1@email.arizona.edu) & Fengqiong Huang (huangfq@emailo.arizona.edu)

#include <mpi.h>
#include <math.h>
#include <stdio.h>
#include <sys/types.h>
#include <stdlib.h>
#include <pthread.h>
#include <sys/time.h>

#define TAG 13
#define PRINT_NODE 3

// helper functions
void print_grid(double **mat, int num_rows, int num_cols) 
{
    int i,j;

    //printf("The %d * %d grid is\n", grid_size, grid_size);
    for(i = 0; i < num_rows; i++)
    {
        for( j = 0; j < num_cols; j++)
        {
          printf("%lf ",  mat[i][j]);
        }
        printf("\n");
    }
}

// Allocate grid
double ** allocate_grid(int row, int col) 
{
    int i;
    double *vals, **temp;

    // allocate values
    vals = (double *) malloc (col * row * sizeof(double));

    // allocate vector of pointers
    temp = (double **) malloc (row * sizeof(double*));

    for(i = 0; i < row; i++)
    {
        temp[i] = &(vals[i * col]);
    }

    return temp;
}

// Exchange edges between workers
void exchange_edges(int w, int num_elements, int height, int num_nodes, double** grid)
{
    if(num_nodes == 1) return;

    // First node (even number) first send messages, then receive messages
    if(w == 0)
    {
        MPI_Send(&grid[height][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD);
        MPI_Recv(&grid[height + 1][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }
    else if(w > 0 && w < num_nodes -1)
    {
        // Even number of node first send messages, then receive messages from odd number of nodes
        if( w % 2 == 0)
        {
            MPI_Send(&grid[1][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD);
            MPI_Send(&grid[height][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD);
            MPI_Recv(&grid[0][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(&grid[height + 1][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
        else
        {
            // Odd number of node first receive messages, then send messages to even number of nodes
            MPI_Recv(&grid[0][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(&grid[height + 1][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Send(&grid[1][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD);
            MPI_Send(&grid[height][0], num_elements, MPI_DOUBLE, w+1, TAG, MPI_COMM_WORLD);
        }
    }
    else if(w == num_nodes - 1)
    {
        // Last node (odd number) first receive messages, then send messages
        MPI_Recv(&grid[0][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Send(&grid[1][0], num_elements, MPI_DOUBLE, w-1, TAG, MPI_COMM_WORLD);
    }

}

int main(int argc, char* argv[])
{
    double spend_time;
    double **grid;
    double old_val, temp;
    int i, j, my_rank, iter, j_start, num_elements;
    int all_initialized; 
    int initialized = 1; 
    double max_diff;
    double my_diff;
    int height;
    double start_time, end_time;
    int num_nodes, grid_size, num_iters;

    MPI_Init(&argc, &argv);

    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &num_nodes);

    grid_size = 1600;
    num_iters = 2;

    // Compute HEIGHT
    height = grid_size / num_nodes;

    grid = allocate_grid(height + 2, grid_size + 2);

    // Initialize the grid 
    int local_grid_size = height + 2; 
    for(i = 0; i < local_grid_size; i++)
    {
        for(j = 0; j < grid_size + 2; j++)
        {
            if(my_rank*height + i == 0 || 
               my_rank*height + i == grid_size + 1 ||
               j == 0 || j == grid_size + 1)
            {
                grid[i][j] = 1.0;
            }
            else
            {
                grid[i][j] = 0.0;
            }
        }
    }
    
    //print for testing
    if(my_rank == PRINT_NODE)
    {
        if(grid_size < 24)
            print_grid(grid, height+2, grid_size+2);
    }

    
    // start timer
    if (my_rank == 0) 
    {
        start_time = MPI_Wtime();
    }

    // Perform red-black grid computation via message passing
    for(iter = 1; iter <= num_iters - 1; iter++)
    {
        // update red points 
        for(i = 1; i <= height; i++)
        {
            if (i % 2 == 1) 
                j_start = 1;
            else 
                j_start = 2;
            for(j = j_start; j <= grid_size; j = j+2)
            {
                grid[i][j] = (grid[i - 1][j] + grid[i][j - 1] + 
                               grid[i + 1][j] + grid[i][j + 1]) * 0.25;
            }
        }

        //print for testing
        if(iter == 29 && my_rank == PRINT_NODE)
        {
            if(grid_size < 24)
                print_grid(grid, height+2, grid_size+2);
        }
        
        // Exchange edges
        exchange_edges(my_rank, grid_size + 2, height, num_nodes, grid);
        
        // update black points
        for(i = 1; i <= height; i++)
        {
            if (i % 2 == 1) 
                j_start = 2;
            else 
                j_start = 1;
            for(j = j_start; j <= grid_size; j = j+2)
            {
                grid[i][j] = (grid[i - 1][j] + grid[i][j - 1] + 
                               grid[i + 1][j] + grid[i][j + 1]) * 0.25;
            }
        }

        // Exchange edges
        exchange_edges(my_rank, grid_size + 2, height, num_nodes, grid);
    }

    // One more updates and keep track of max_diff
    // update red points 
    for(i = 1; i <= height; i++)
    {
        if (i % 2 == 1) 
            j_start = 1;
        else 
            j_start = 2;
        for(j = j_start; j <= grid_size; j = j+2)
        {
            old_val = grid[i][j];
            grid[i][j] = (grid[i - 1][j] + grid[i][j - 1] + 
                           grid[i + 1][j] + grid[i][j + 1]) * 0.25;
            temp = old_val - grid[i][j];
            if(temp < 0)
                temp = -temp;
            if(temp > my_diff)
                my_diff = temp;
        }
    }

    // Exchange edges
    exchange_edges(my_rank, grid_size + 2, height, num_nodes, grid);

    // update black points
    for(i = 1; i <= height; i++)
    {
        if (i % 2 == 1) 
            j_start = 2;
        else 
            j_start = 1;
        for(j = j_start; j <= grid_size; j = j+2)
        {
            old_val = grid[i][j];
            grid[i][j] = (grid[i - 1][j] + grid[i][j - 1] + 
                           grid[i + 1][j] + grid[i][j + 1]) * 0.25;
            temp = old_val - grid[i][j];
            if(temp < 0)
                temp = -temp;
            if(temp > my_diff)
                my_diff = temp;
        }
    }

    // update the max_diff
    MPI_Allreduce(&my_diff, &max_diff, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);

    // stop timer
    if(my_rank == 0) 
    {
        end_time = MPI_Wtime();
        spend_time = end_time - start_time;
        printf("Maximum diff is: %f\n", max_diff);
        double tt = end_time - start_time;
        printf("Compute time for grid-size(%d), iterations(%d), nodes(%d)"
                "is: %.3f seconds\n", grid_size, num_iters, num_nodes, tt);
    }

    free(grid);
    MPI_Finalize();

    exit(1);
}

